import os
import sys
import json
import multiprocessing
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import argparse
import time

# 导入你的预处理器模块
from preprocessor import process_package

def process_single_package(package_info):
    """
    处理单个包的包装函数，用于多进程处理
    
    Args:
        package_info (tuple): (package_dir, package_name, package_type)
    
    Returns:
        dict: 处理结果
    """
    package_dir, package_name, package_type = package_info
    
    try:
        # 调用你的预处理器
        js_files = process_package(package_dir)
        
        return {
            'package_name': package_name,
            'package_type': package_type,
            'package_dir': package_dir,
            'status': 'success',
            'js_files_count': len(js_files),
            'js_files': js_files,
            'has_virtual_file': any('_virtual_package_script.js' in f for f in js_files),
            'error': None
        }
    
    except Exception as e:
        return {
            'package_name': package_name,
            'package_type': package_type,
            'package_dir': package_dir,
            'status': 'error',
            'js_files_count': 0,
            'js_files': [],
            'has_virtual_file': False,
            'error': str(e)
        }

def collect_packages(base_dir):
    """
    收集所有需要处理的包
    
    Args:
        base_dir (str): 基础目录路径
    
    Returns:
        list: 包信息列表
    """
    packages = []
    base_path = Path(base_dir)
    
    # 遍历 malicious 和 benign 目录
    for package_type in ['malicious', 'benign']:
        type_dir = base_path / package_type
        
        if not type_dir.exists():
            print(f"⚠️  目录不存在: {type_dir}")
            continue
        
        print(f"🔍 收集 {package_type} 包...")
        
        for package_dir in type_dir.iterdir():
            if package_dir.is_dir():
                packages.append((
                    str(package_dir),
                    package_dir.name,
                    package_type
                ))
    
    return packages

def save_results(results, output_file):
    """
    保存处理结果到JSON文件
    """
    # 统计信息
    total_packages = len(results)
    successful_packages = len([r for r in results if r['status'] == 'success'])
    failed_packages = len([r for r in results if r['status'] == 'error'])
    packages_with_virtual_files = len([r for r in results if r['has_virtual_file']])
    
    malicious_packages = [r for r in results if r['package_type'] == 'malicious']
    benign_packages = [r for r in results if r['package_type'] == 'benign']
    
    malicious_with_virtual = len([r for r in malicious_packages if r['has_virtual_file']])
    benign_with_virtual = len([r for r in benign_packages if r['has_virtual_file']])
    
    summary = {
        'processing_summary': {
            'total_packages': total_packages,
            'successful_packages': successful_packages,
            'failed_packages': failed_packages,
            'success_rate': round(successful_packages / total_packages * 100, 2) if total_packages > 0 else 0,
            'packages_with_virtual_files': packages_with_virtual_files,
            'virtual_file_rate': round(packages_with_virtual_files / total_packages * 100, 2) if total_packages > 0 else 0
        },
        'by_type': {
            'malicious': {
                'total': len(malicious_packages),
                'with_virtual_files': malicious_with_virtual,
                'virtual_file_rate': round(malicious_with_virtual / len(malicious_packages) * 100, 2) if malicious_packages else 0
            },
            'benign': {
                'total': len(benign_packages),
                'with_virtual_files': benign_with_virtual,
                'virtual_file_rate': round(benign_with_virtual / len(benign_packages) * 100, 2) if benign_packages else 0
            }
        },
        'detailed_results': results
    }
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)

def print_summary(results):
    """
    打印处理摘要
    """
    total = len(results)
    successful = len([r for r in results if r['status'] == 'success'])
    failed = len([r for r in results if r['status'] == 'error'])
    with_virtual = len([r for r in results if r['has_virtual_file']])
    
    malicious = [r for r in results if r['package_type'] == 'malicious']
    benign = [r for r in results if r['package_type'] == 'benign']
    
    malicious_virtual = len([r for r in malicious if r['has_virtual_file']])
    benign_virtual = len([r for r in benign if r['has_virtual_file']])
    
    print("\n" + "=" * 60)
    print("📊 批量预处理完成摘要")
    print("=" * 60)
    
    if total == 0:
        print("❌ 没有处理任何包")
        return
    
    print(f"📦 总包数: {total}")
    print(f"✅ 成功处理: {successful} ({successful/total*100:.1f}%)")
    print(f"❌ 处理失败: {failed} ({failed/total*100:.1f}%)")
    print(f"🔧 生成虚拟文件: {with_virtual} ({with_virtual/total*100:.1f}%)")
    
    print(f"\n📈 按类型统计:")
    
    # 安全计算恶意包百分比
    if len(malicious) > 0:
        malicious_percent = malicious_virtual/len(malicious)*100
        print(f"  🦠 恶意包: {len(malicious)} 个，其中 {malicious_virtual} 个生成了虚拟文件 ({malicious_percent:.1f}%)")
    else:
        print(f"  🦠 恶意包: 0 个")
    
    # 安全计算良性包百分比
    if len(benign) > 0:
        benign_percent = benign_virtual/len(benign)*100
        print(f"  ✨ 良性包: {len(benign)} 个，其中 {benign_virtual} 个生成了虚拟文件 ({benign_percent:.1f}%)")
    else:
        print(f"  ✨ 良性包: 0 个")
    
    # 显示一些失败的案例
    failed_results = [r for r in results if r['status'] == 'error']
    if failed_results:
        print(f"\n❌ 失败案例 (前5个):")
        for i, result in enumerate(failed_results[:5], 1):
            print(f"  {i}. {result['package_name']}: {result['error']}")
    
    # 如果没有恶意包或良性包，给出提示
    if len(malicious) == 0 and len(benign) == 0:
        print(f"\n⚠️  注意: 没有找到标准的 'malicious' 或 'benign' 目录")
        print(f"   如果你处理的是未标记的真实数据，这是正常的")
        
        # 显示实际找到的包类型
        unique_types = set([r['package_type'] for r in results])
        if unique_types:
            print(f"   实际包类型: {', '.join(unique_types)}")
def main():
    parser = argparse.ArgumentParser(description='批量NPM包预处理器')
    parser.add_argument('--base_dir', 
                       default=r'F:\robust_extracted_packages',
                       help='包含malicious和benign子目录的基础目录')
    parser.add_argument('--output', 
                       default='batch_preprocessing_results.json',
                       help='输出结果文件路径')
    parser.add_argument('--workers', 
                       type=int, 
                       default=multiprocessing.cpu_count(),
                       help='并行处理的进程数')
    parser.add_argument('--sample', 
                       type=int, 
                       help='仅处理指定数量的包（用于测试）')
    
    args = parser.parse_args()
    
    print("🚀 开始批量预处理NPM包...")
    print(f"📁 基础目录: {args.base_dir}")
    print(f"💾 输出文件: {args.output}")
    print(f"⚡ 并行进程数: {args.workers}")
    
    # 收集所有包
    start_time = time.time()
    packages = collect_packages(args.base_dir)
    
    if not packages:
        print("❌ 未找到任何包！请检查目录路径。")
        return
    
    # 如果指定了sample参数，只处理部分包
    if args.sample:
        packages = packages[:args.sample]
        print(f"🧪 测试模式：仅处理前 {len(packages)} 个包")
    
    print(f"📊 发现 {len(packages)} 个包待处理")
    
    # 使用多进程并行处理
    results = []
    
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        # 提交所有任务
        future_to_package = {
            executor.submit(process_single_package, pkg_info): pkg_info 
            for pkg_info in packages
        }
        
        # 使用进度条显示处理进度
        with tqdm(total=len(packages), desc="处理包", unit="pkg") as pbar:
            for future in as_completed(future_to_package):
                try:
                    result = future.result()
                    results.append(result)
                    
                    # 更新进度条描述
                    if result['status'] == 'success':
                        pbar.set_postfix(status="✅", package=result['package_name'][:20])
                    else:
                        pbar.set_postfix(status="❌", package=result['package_name'][:20])
                    
                except Exception as e:
                    # 处理未预期的错误
                    pkg_info = future_to_package[future]
                    results.append({
                        'package_name': pkg_info[1],
                        'package_type': pkg_info[2],
                        'package_dir': pkg_info[0],
                        'status': 'error',
                        'js_files_count': 0,
                        'js_files': [],
                        'has_virtual_file': False,
                        'error': f"Unexpected error: {str(e)}"
                    })
                    pbar.set_postfix(status="💥", package=pkg_info[1][:20])
                
                pbar.update(1)
    
    # 保存结果
    save_results(results, args.output)
    
    # 打印摘要
    print_summary(results)
    
    elapsed_time = time.time() - start_time
    print(f"\n⏱️  总用时: {elapsed_time:.2f} 秒")
    print(f"💾 详细结果已保存到: {args.output}")

if __name__ == "__main__":
    main()